# -*- coding: utf-8 -*-
# @Time    : 2020/12/4 9:36
# @Author  : DuJiabao
# @File    : BlackgroundGenerator.py
# @Description : 这是一个用于生成背景的文件

import numpy as np
import cv2


class SingleColorRangeHSV(object):
    """
    这是一个存储HSV色彩范围的类
    """
    def __init__(self, hLow, sLow, vLow, hHigh, sHigh, vHigh):
        """
        初始化函数，初始化该类的色彩范围
        :param hLow:
        :param sLow:
        :param vLow:
        :param hHigh:
        :param sHigh:
        :param vHigh:
        """
        self.hLow = hLow
        self.sLow = sLow
        self.vLow = vLow
        self.hHigh = hHigh
        self.sHigh = sHigh
        self.vHigh = vHigh

    def getLow(self):
        return self.hLow, self.sLow, self.vLow

    def getHigh(self):
        return self.hHigh, self.sHigh, self.vHigh

    def getH(self):
        return self.hLow, self.hHigh

    def getS(self):
        return self.sLow, self.sHigh

    def getV(self):
        return self.vLow, self.vHigh


class ColorRangeHSV(object):
    """
    存储常用颜色的色彩范围
    """
    def __init__(self):
        self.blackRange = SingleColorRangeHSV(0, 0, 0, 180, 255, 46)
        self.whiteRange = SingleColorRangeHSV(0, 0, 221, 180, 30, 255)
        self.grayRange = SingleColorRangeHSV(0, 0, 46, 180, 43, 220)
        self.orangeRange = SingleColorRangeHSV(11, 43, 46, 25, 255, 255)
        self.yellowRange = SingleColorRangeHSV(26, 43, 46, 34, 255, 255)
        self.greenRange = SingleColorRangeHSV(35, 43, 46, 77, 255, 255)
        self.blueRange = SingleColorRangeHSV(100, 43, 46, 124, 255, 255)
        self.purpleRange = SingleColorRangeHSV(125, 43, 46, 155, 255, 255)
        self.redRange = [SingleColorRangeHSV(0, 10, 43, 255, 46, 255),
                         SingleColorRangeHSV(156, 180, 43, 255, 46, 255)]


def generateBlackground(shape, color_range: SingleColorRangeHSV, isPure=False):
    """
    生成背景，自定义颜色随机的范围
    :param isPure: 是否采用纯净的背景
    :param shape: 背景形状: 允许长度为[1,2,3]的tuple、list; int、float类型
    :param color_range: 传入SingleColorRangeHSV类。
    :return: BGR格式的图片
    """

    if type(shape) in [tuple, list]:
        if len(shape) == 2:
            shape = (shape[0], shape[1], 3)
        elif len(shape) == 1:
            shape = (shape[0], shape[0], 3)
    elif type(shape) in [int, float]:
        shape = (int(shape), int(shape), 3)

    bg = np.zeros(shape, dtype=np.uint8)

    if isPure:
        randomShape = 1
    else:
        randomShape = list(shape)[:-1]
    bg[:, :, 0] = np.random.randint(color_range.getH()[0], color_range.getH()[1], randomShape)
    bg[:, :, 1] = np.random.randint(color_range.getS()[0], color_range.getS()[1], randomShape)
    bg[:, :, 2] = np.random.randint(color_range.getV()[0], color_range.getV()[1], randomShape)

    bg = cv2.cvtColor(bg, cv2.COLOR_HSV2BGR)
    return bg


if __name__ == '__main__':
    colorRange = ColorRangeHSV()
    bg = generateBlackground(500, colorRange.yellowRange, True)
    cv2.imshow("bg", bg)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
